#include <windows.h>
#include <tlhelp32.h>
#include <iostream>
#include <string>
#include <tchar.h>
#include <commctrl.h>
#include <sddl.h>


#pragma comment(lib, "comctl32.lib")

#define ID_BUTTON_INJECT 1

HWND hwndProcessName, hwndDllPath;

DWORD GetProcessIdByName(const wchar_t* processName) {
    PROCESSENTRY32W processEntry;
    processEntry.dwSize = sizeof(PROCESSENTRY32W);

    HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (snapshot == INVALID_HANDLE_VALUE) {
        return 0;
    }

    if (Process32FirstW(snapshot, &processEntry)) {
        do {
            if (_wcsicmp(processEntry.szExeFile, processName) == 0) {
                CloseHandle(snapshot);
                return processEntry.th32ProcessID;
            }
        } while (Process32NextW(snapshot, &processEntry));
    }

    CloseHandle(snapshot);
    return 0;
}

bool InjectDLL(DWORD processId, const wchar_t* dllPath) {
    HANDLE processHandle = OpenProcess(PROCESS_ALL_ACCESS, FALSE, processId);
    if (!processHandle) {
        MessageBox(NULL, L"无法打开目标进程。", L"错误", MB_OK | MB_ICONERROR);
        return false;
    }

    size_t dllPathSize = (wcslen(dllPath) + 1) * sizeof(wchar_t);

    LPVOID remoteMemory = VirtualAllocEx(processHandle, nullptr, dllPathSize, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
    if (!remoteMemory) {
        MessageBox(NULL, L"无法在目标进程中分配内存。", L"错误", MB_OK | MB_ICONERROR);
        CloseHandle(processHandle);
        return false;
    }

    if (!WriteProcessMemory(processHandle, remoteMemory, dllPath, dllPathSize, nullptr)) {
        MessageBox(NULL, L"无法写入DLL路径到目标进程。", L"错误", MB_OK | MB_ICONERROR);
        VirtualFreeEx(processHandle, remoteMemory, 0, MEM_RELEASE);
        CloseHandle(processHandle);
        return false;
    }

    LPVOID loadLibraryAddress = (LPVOID)GetProcAddress(GetModuleHandleW(L"kernel32.dll"), "LoadLibraryW");
    if (!loadLibraryAddress) {
        MessageBox(NULL, L"无法获取LoadLibraryW函数的地址。", L"错误", MB_OK | MB_ICONERROR);
        VirtualFreeEx(processHandle, remoteMemory, 0, MEM_RELEASE);
        CloseHandle(processHandle);
        return false;
    }

    HANDLE remoteThread = CreateRemoteThread(processHandle, nullptr, 0, (LPTHREAD_START_ROUTINE)loadLibraryAddress, remoteMemory, 0, nullptr);
    if (!remoteThread) {
        MessageBox(NULL, L"无法创建远程线程。", L"错误", MB_OK | MB_ICONERROR);
        VirtualFreeEx(processHandle, remoteMemory, 0, MEM_RELEASE);
        CloseHandle(processHandle);
        return false;
    }

    WaitForSingleObject(remoteThread, INFINITE);

    VirtualFreeEx(processHandle, remoteMemory, 0, MEM_RELEASE);
    CloseHandle(remoteThread);
    CloseHandle(processHandle);

    return true;
}

void RequestAdminPrivileges() {
    SHELLEXECUTEINFOW sei = { sizeof(sei) };
    sei.lpVerb = L"runas";
    sei.lpFile = GetCommandLineW();
    sei.hwnd = nullptr;
    sei.nShow = SW_SHOWNORMAL;

    if (!ShellExecuteExW(&sei)) {
        MessageBox(nullptr, L"需要管理员权限才能运行此程序。", L"错误", MB_OK | MB_ICONERROR);
        exit(1);
    }
}

void AutoInjectGTA5() {
    const wchar_t* processName = L"GTA5.exe";
    wchar_t dllPath[MAX_PATH] = L"YimMenu.dll";

    // 检查当前目录下是否有 YimMenu.dll
    if (GetFileAttributesW(dllPath) == INVALID_FILE_ATTRIBUTES) {
        // 弹出文件选择框
        OPENFILENAMEW ofn;
        wchar_t szFile[260] = { 0 };
        ZeroMemory(&ofn, sizeof(ofn));
        ofn.lStructSize = sizeof(ofn);
        ofn.hwndOwner = nullptr;
        ofn.lpstrFile = szFile;
        ofn.nMaxFile = sizeof(szFile) / sizeof(wchar_t);
        ofn.lpstrFilter = L"DLL Files\0*.dll\0All Files\0*.*\0";
        ofn.nFilterIndex = 1;
        ofn.lpstrFileTitle = nullptr;
        ofn.nMaxFileTitle = 0;
        ofn.lpstrInitialDir = nullptr;
        ofn.Flags = OFN_PATHMUSTEXIST | OFN_FILEMUSTEXIST;

        if (GetOpenFileNameW(&ofn) == TRUE) {
            wcscpy_s(dllPath, szFile);
        } else {
            MessageBoxW(nullptr, L"未找到YimMenu.dll文件或未选择DLL文件。", L"错误", MB_OK | MB_ICONERROR);
            return;
        }
    } else {
        // 获取当前工作目录并构造完整 DLL 路径
        wchar_t currentDir[MAX_PATH];
        GetCurrentDirectoryW(MAX_PATH, currentDir);
        wcscat_s(currentDir, L"\\YimMenu.dll");
        wcscpy_s(dllPath, currentDir);
    }

    DWORD processId = GetProcessIdByName(processName);
    if (processId == 0) {
        MessageBoxW(nullptr, L"未找到GTA5进程，请确保GTA5正在运行。", L"错误", MB_OK | MB_ICONERROR);
        return;
    }

    if (!InjectDLL(processId, dllPath)) {
        MessageBoxW(nullptr, L"自动注入失败，请手动尝试注入。", L"错误", MB_OK | MB_ICONERROR);
    }
}


LRESULT CALLBACK WindowProc(HWND hwnd, UINT uMsg, WPARAM wParam, LPARAM lParam) {
    switch (uMsg) {
        case WM_CREATE: {
            CreateWindow(L"STATIC", L"目标进程名称:", WS_VISIBLE | WS_CHILD, 20, 20, 120, 25, hwnd, NULL, NULL, NULL);
            hwndProcessName = CreateWindow(L"EDIT", L"GTA5.exe", WS_VISIBLE | WS_CHILD | WS_BORDER, 150, 20, 200, 25, hwnd, NULL, NULL, NULL);

            CreateWindow(L"STATIC", L"DLL 路径:", WS_VISIBLE | WS_CHILD, 20, 60, 120, 25, hwnd, NULL, NULL, NULL);
            hwndDllPath = CreateWindow(L"EDIT", L"YimMenu.dll", WS_VISIBLE | WS_CHILD | WS_BORDER, 150, 60, 200, 25, hwnd, NULL, NULL, NULL);

            CreateWindow(L"BUTTON", L"注入DLL", WS_VISIBLE | WS_CHILD, 150, 100, 100, 30, hwnd, (HMENU)ID_BUTTON_INJECT, NULL, NULL);
            break;
        }
        case WM_COMMAND: {
            if (LOWORD(wParam) == ID_BUTTON_INJECT) {
                wchar_t processName[256];
                wchar_t dllPath[256];

                GetWindowTextW(hwndProcessName, processName, sizeof(processName) / sizeof(wchar_t));
                GetWindowTextW(hwndDllPath, dllPath, sizeof(dllPath) / sizeof(wchar_t));

                DWORD processId = GetProcessIdByName(processName);
                if (processId == 0) {
                    MessageBox(hwnd, L"未找到目标进程。", L"错误", MB_OK | MB_ICONERROR);
                    return 0;
                }

                if (InjectDLL(processId, dllPath)) {
                    MessageBox(hwnd, L"DLL 注入成功！", L"成功", MB_OK | MB_ICONINFORMATION);
                } else {
                    MessageBox(hwnd, L"DLL 注入失败。", L"失败", MB_OK | MB_ICONERROR);
                }
            }
            break;
        }
        case WM_DESTROY:
            PostQuitMessage(0);
            break;
        default:
            return DefWindowProc(hwnd, uMsg, wParam, lParam);
    }
    return 0;
}

bool IsUserAdmin() {
    BOOL isAdmin = FALSE;
    PSID administratorsGroup = NULL;

    SID_IDENTIFIER_AUTHORITY ntAuthority = SECURITY_NT_AUTHORITY;
    if (AllocateAndInitializeSid(&ntAuthority, 2,
                                 SECURITY_BUILTIN_DOMAIN_RID,
                                 DOMAIN_ALIAS_RID_ADMINS,
                                 0, 0, 0, 0, 0, 0,
                                 &administratorsGroup)) {
        if (!CheckTokenMembership(NULL, administratorsGroup, &isAdmin)) {
            isAdmin = FALSE;
        }
        FreeSid(administratorsGroup);
    }

    return isAdmin == TRUE;
}

int main(int argc, char *argv[]) {
    HINSTANCE hInstance = GetModuleHandle(NULL);
    int nCmdShow = SW_SHOWNORMAL;

    // 将原有的 wWinMain 的逻辑放到 main 中
    if (!IsUserAdmin()) {
        RequestAdminPrivileges();
    }

    AutoInjectGTA5();

    const wchar_t CLASS_NAME[] = L"DLL Injector Window Class";

    WNDCLASS wc = {};
    wc.lpfnWndProc = WindowProc;
    wc.hInstance = hInstance;
    wc.lpszClassName = CLASS_NAME;
    wc.hbrBackground = (HBRUSH)(COLOR_WINDOW + 1);

    RegisterClass(&wc);

    HWND hwnd = CreateWindowEx(
            0,
            CLASS_NAME,
            L"DLL 注入器",
            WS_OVERLAPPEDWINDOW,
            CW_USEDEFAULT, CW_USEDEFAULT, 400, 200,
            nullptr,
            nullptr,
            hInstance,
            nullptr
    );

    if (hwnd == nullptr) {
        return 0;
    }

    ShowWindow(hwnd, nCmdShow);

    MSG msg = {};
    while (GetMessage(&msg, nullptr, 0, 0)) {
        TranslateMessage(&msg);
        DispatchMessage(&msg);
    }

    return 0;
}